# from argparse import ArgumentParser
import time

from PIL import Image, ImageFont, ImageDraw
from tensorflow.python.tools import optimize_for_inference_lib

from combine import input_path, gen_frozen_graph
from config import Input_shape, channels, score_threshold, ignore_thresh, path, data_path, model_type, MODEL_PATH, \
    MAX_BOX_NUM
from detect_function import predict
from utils.yolo_utils import read_anchors, letterbox_image  # , resize_image
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import tensorflow as tf
import argparse
import os

from yolo_v3 import yolo_v3
from yolo_v3_tiny import yolo_v3_tiny

os.environ["CUDA_VISIBLE_DEVICES"] = "0"


# from tensorflow.python.client import device_lib
# print(device_lib.list_local_devices())


class YOLO(object):
    def __init__(self, time_dir, model_index=None):
        
        # self.anchors_path = path + '/yolo3/model/yolo_anchors.txt'
        if model_type == 'N':
            self.anchors_path = data_path + '/normal_anchor.txt'
        else:
            self.anchors_path = data_path + '/tiny_anchor.txt'
        
        self.trainable = True
        self.time_dir = time_dir
        self.model_index = model_index
        # args = self.argument()
        # if args.COCO:
        #     print("-----------COCO-----------")
        #     self.COCO = True
        #     self.classes_path = self.PATH + '/model/coco_classes.txt'
        #     self.trainable = False
        # elif args.VOC:
        #     print("-----------VOC------------")
        #     self.classes_path = self.PATH + '/model/voc_classes.txt'
        # elif args.boat:
        #     print("-----------boat-----------")
        #     self.classes_path = self.PATH + '/model/boat_classes.txt'
        
        self.anchors = read_anchors(self.anchors_path)
        self.threshold = score_threshold  # threshold
        self.ignore_thresh = ignore_thresh
        self.INPUT_SIZE = (Input_shape, Input_shape)  # fixed size or (None, None)
        self.is_fixed_size = self.INPUT_SIZE != (None, None)
        # LOADING SESSION...    [M, 4] [M, 1]
        self.test_list = []
        self.boxes, self.scores, self.seeds, self.classes, self.sess, self.var_ckpt = self.load()
    
    @staticmethod
    def argument():
        parser = argparse.ArgumentParser(description='COCO or VOC or boat')
        parser.add_argument('--COCO', action='store_true', help='COCO flag')
        parser.add_argument('--VOC', action='store_true', help='VOC flag')
        parser.add_argument('--boat', action='store_true', help='boat flag')
        args = parser.parse_args()
        return args
    
    def load(self):
        # Remove nodes from graph or reset entire default graph
        tf.reset_default_graph()


        # Generate output tensor targets for filtered bounding boxes.
        self.x = tf.placeholder(tf.float32, shape=[None, Input_shape, Input_shape, channels],name='input')
        self.image_shape = tf.placeholder(tf.float32, shape=[2,],name='image_shape')
        self.phase = tf.placeholder(tf.bool, shape=[], name='phase')
        
        # self.is_training = tf.placeholder(tf.bool)
        # image_shape = np.array([image.size[0], image.size[1]])  # tf.placeholder(tf.float32, shape=[2,])
        
        # Generate output tensor targets for filtered bounding boxes.
        if model_type == 'N':
            model = yolo_v3
        else:
            model = yolo_v3_tiny
        scale_total = model(self.x, self.phase)
        
        # detect  [M, 4] 上左下右
        boxes, scores, seeds, classes = predict(scale_total, self.anchors, self.image_shape,
                                                max_boxes=MAX_BOX_NUM, score_threshold=self.threshold,
                                                iou_threshold=self.ignore_thresh)
        classes = tf.expand_dims(tf.cast(classes, tf.float32), axis=-1)
        output_node = tf.concat([boxes, seeds, classes], axis=-1, name="output_node")
        # boxes = tf.identity(boxes,name="output_node")
        # Add ops to save and restore all the variables
        var_list = tf.trainable_variables()
        g_list = tf.global_variables()
        
        ########## 确认BN 模块中的名字是否如下，如不一致，将不会保存！！！！
        bn_moving_vars = [g for g in g_list if 'moving_avg' in g.name or 'moving_var' in g.name]
        # bn_moving_vars += [g for g in g_list if 'moving_var' in g.name]
        var_list += bn_moving_vars
        
        #############
        saver = tf.train.Saver(var_list)
        
        # Allowing GPU memory growth
        
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        sess = tf.Session(config=config)
        # sess.run(tf.global_variables_initializer())
        
        load_checkpoints_dir = os.path.join(MODEL_PATH, self.time_dir)
        # var_file = tf.train.latest_checkpoint(load_checkpoints_dir)
        if self.model_index is None:
            var_file = tf.train.latest_checkpoint(load_checkpoints_dir)
        else:
            var_file = os.path.join(load_checkpoints_dir, 'model.ckpt-' + self.model_index)
        
        saver.restore(sess, var_file)  # 从模型中恢复最新变量
        
        return boxes, scores, seeds, classes, sess, var_file
    
    def detect_image(self, image, filename, need_vis=True, save_path=None):
        
        if self.is_fixed_size:
            assert self.INPUT_SIZE[0] % 32 == 0, 'Multiples of 32 required'
            assert self.INPUT_SIZE[1] % 32 == 0, 'Multiples of 32 required'
            boxed_image, image_shape = letterbox_image(image, tuple(reversed(self.INPUT_SIZE)))
            # boxed_image, image_shape = resize_image(image, tuple(reversed(self.INPUT_SIZE)))
        else:
            new_image_size = (image.width - (image.width % 32), image.height - (image.height % 32))
            boxed_image, image_shape = letterbox_image(image, new_image_size)
            # boxed_image, image_shape = resize_image(image, new_image_size)
        image_data = np.array(boxed_image, dtype='float32')
        
        print("heights, widths:", image_shape)
        image_data /= 255.
        inputs = np.expand_dims(image_data, 0)  # Add batch dimension. #
        time_start = time.time()
        
        # [M, 4]  上左下右
        out_boxes, out_scores, out_seeds, out_classes = self.sess.run(
            [self.boxes, self.scores, self.seeds, self.classes],
            feed_dict={self.x: inputs,
                       self.image_shape: image_shape,
                       self.phase: False
                       # self.is_training: False
                       })
        time_end = time.time()
        print('time cost', time_end - time_start, 's')
        
        print('\nFound {} boxes for {}'.format(len(out_boxes), 'img'))
        
        # Visualisation#################################################################################################
        if need_vis:
            fig, ax = plt.subplots(1)
            ax.imshow(image)
            if filename is not None:
                ax.annotate(filename, (20, 40), color='b',
                            weight='bold',
                            fontsize=9, ha='left', va='center')
            
            for i in range(out_boxes.shape[0]):
                # 原图中的大小
                box = out_boxes[i]
                seed = out_seeds[i]
                # print(box)
                # 数据制作的标准： 左上右下, 数据输出格式：上左下右
                top, left, bottom, right = box  # y_min, x_min, y_max, x_max
                # left,top,  right, bottom = box  # y_min, x_min, y_max, x_max
                top = max(0, np.floor(top + 0.5).astype(np.int32))
                left = max(0, np.floor(left + 0.5).astype(np.int32))
                bottom = min(image.size[1], np.floor(bottom + 0.5).astype(np.int32))
                right = min(image.size[0], np.floor(right + 0.5).astype(np.int32))
                width = right - left
                height = bottom - top
                
                # My kingdom for a good redistributable image drawing library.
                rect = patches.Rectangle((left, top), width, height, linewidth=1, edgecolor='r', facecolor='none')
                ax.add_patch(rect)
                
                # # 显示 分类
                ax.annotate('%d' % out_classes[i], (left + width / 5, top + height / 5), weight='bold',
                            fontsize=11, ha='center', va='center')
                # # 显示 顺序
                # ax.annotate('%d' % i, (left+width/5, top+height*4/5),  weight='bold',
                #             fontsize=11, ha='center', va='center')
                
                # 显示 概率
                ax.annotate('%.1f' % out_scores[i], (left + width * 4 / 5, top + height * 4 / 5),  # color='b',
                            fontsize=8, ha='center', va='center')
                
                plt.plot([seed[1]], [seed[0]], "ro")
            
            if save_path is not None:
                plt.savefig(save_path, format='png')
            else:
                plt.show()
                plt.close()
    
    def write_pb(self, need_optimize=False):
        gd = self.sess.graph.as_graph_def()
        
        for node in gd.node:
            if node.op == 'RefSwitch':
                node.op = 'Switch'
                for index in range(len(node.input)):
                    if 'moving_' in node.input[index]:
                        node.input[index] = node.input[index] + '/read'
            
            elif node.op == 'AssignSub':
                node.op = 'Sub'
                if 'use_locking' in node.attr:
                    del node.attr['use_locking']
            elif node.op == 'AssignAdd':
                node.op = 'Add'
        load_checkpoints_dir = os.path.join(MODEL_PATH, self.time_dir)
        tf.train.write_graph(gd, load_checkpoints_dir, input_path)
        gen_frozen_graph(self.var_ckpt, load_checkpoints_dir)
        
        if need_optimize:
            input_graph_def = tf.GraphDef()
            with tf.gfile.Open(os.path.join(load_checkpoints_dir, 'output_graph.pb'), "rb") as f:
                data = f.read()
            input_graph_def.ParseFromString(data)
            
            output_graph_def = optimize_for_inference_lib.optimize_for_inference(
                input_graph_def,
                ["input", "image_shape", "phase"],  # an array of the input node(s)
                ["output_node"],  # an array of output nodes
                # tf.float32.as_datatype_enum
                [tf.float32.as_datatype_enum, tf.float32.as_datatype_enum, tf.bool.as_datatype_enum]
            )
            
            # Save the optimized graph
            
            f = tf.gfile.FastGFile(os.path.join(load_checkpoints_dir, 'optimized_graph.pb'), "w")
            f.write(output_graph_def.SerializeToString())


def detect_img(yolo, annotation_path, need_vis, save_root):
    # data_path='F:/tf_projects/data/VOCdevkit/VOC2012/JPEGImages/'
    # annotation_path='F:/tf_projects/CNN/yolo3/model/voc_train.txt'
    if save_root is not None:
        if not os.path.exists(save_root):
            os.mkdir(save_root)
    if os.path.isfile(annotation_path):
        with open(annotation_path) as f:
            GG = f.readlines()
            # np.random.shuffle(GG)
            for i, line in enumerate(GG):
                # input('next')
                if line == '\n':
                    continue
                if line.startswith('#'):
                    continue
                
                line = line.split('.png')
                filepath = line[0] + '.png'
                filename = line[0].split('/')[-3] + '%s.png' % i
                image = Image.open(filepath)
                if save_root is None:
                    save_path = None
                else:
                    save_path = save_root + '/' + filename
                yolo.detect_image(image, filename, need_vis=need_vis, save_path=save_path)
    elif os.path.isdir(annotation_path):
        for filename in os.listdir(annotation_path):
            filepath = os.path.join(annotation_path, filename)
            if os.path.isfile(filepath):
                image = Image.open(filepath)
                if save_root is None:
                    save_path = None
                else:
                    save_path = save_root + '/' + filename
                yolo.detect_image(image, filename, need_vis, save_path)


def write_graph(yolo):
    yolo.write_pb()


if __name__ == '__main__':
    # write_graph(YOLO('20190626-2018/rutine', '15'))
    
    annotation_path = 'F:/ProjectData/yolo3/test/testpic/problem'
    # save_root = 'F:/ProjectData/yolo3/test/testpic/yuxiaoye'
    detect_img(YOLO('20190626-2018/rutine', '15'), annotation_path, need_vis=True, save_root=None)  # iou 0.5  best
    # #
    # #
